from tqsdk import TqApi, TqAuth, TqKq
from tqsdk.ta import BOLL
import time

def boll_strategy():
    # 创建天勤API实例
    api = TqApi(auth=TqAuth("cps168", "alibaba999"), account=TqKq())
    
    # 获取K线数据(螺纹钢主力合约)
    klines = api.get_kline_serial("SHFE.rb2505", 900, 300)  # 15分钟K线
    
    # 持仓变量
    position = None
    
    try:
        while True:
            # 同步等待K线更新
            deadline = time.time() + 1
            api.wait_update(deadline=deadline)
            
            # 确保有足够K线数据(布林带通常需要20根K线)
            if len(klines.close) < 20:
                print("等待K线数据积累...")
                continue
                
            try:
                # 计算布林带指标
                boll_data = BOLL(klines, 20, 2)  # 20周期，2倍标准差
                
                # 获取布林带值
                upper = boll_data["upper"][-1]
                lower = boll_data["lower"][-1]
                middle = boll_data["mid"][-1]
                close = klines.close[-1]
                
            except Exception as e:
                print(f"指标计算错误: {e}")
                continue
            
            # 获取当前持仓
            position = api.get_position("SHFE.rb2505")
            
            # 交易逻辑
            if close >= upper and (position is None or position.pos == 0):
                print("价格触及上轨，开空仓")
                api.insert_order(symbol="SHFE.rb2505", direction="SELL", offset="OPEN", volume=1)
                
            elif close <= lower and (position is not None and position.pos < 0):
                print("价格触及下轨，平空仓")
                api.insert_order(symbol="SHFE.rb2505", direction="BUY", offset="CLOSE", volume=1)
                
            elif close <= lower and (position is None or position.pos == 0):
                print("价格触及下轨，开多仓")
                api.insert_order(symbol="SHFE.rb2505", direction="BUY", offset="OPEN", volume=1)
                
            elif close >= middle and (position is not None and position.pos > 0):
                print("价格回升至中轨，平多仓")
                api.insert_order(symbol="SHFE.rb2505", direction="SELL", offset="CLOSE", volume=1)
                
    finally:
        api.close()

if __name__ == "__main__":
    boll_strategy()